import os
import librosa
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
import librosa.display
import torch.nn as nn
import torch.optim as optim
from sklearn.metrics import average_precision_score


# 加载音频并统一采样率和时长
def load_audio(filepath, target_sr=16000, max_duration=10):
    audio, sr = librosa.load(filepath, sr=target_sr)
    if len(audio) > target_sr * max_duration:
        audio = audio[:target_sr * max_duration]
    else:
        padding = target_sr * max_duration - len(audio)
        audio = np.pad(audio, (0, padding))
    return audio


# 提取梅尔频谱特征
def extract_mel_spectrogram(audio, sr=16000, n_mels=128, n_fft=2048, hop_length=512):
    mel_spec = librosa.feature.melspectrogram(y=audio, sr=sr, n_mels=n_mels,
                                              n_fft=n_fft, hop_length=hop_length)
    log_mel_spec = librosa.power_to_db(mel_spec, ref=np.max)
    return log_mel_spec


# 自定义数据集
class BirdDataset(Dataset):
    def __init__(self, file_paths, labels, sr=16000, max_duration=10, n_mels=128):
        self.file_paths = file_paths
        self.labels = labels
        self.sr = sr
        self.max_duration = max_duration
        self.n_mels = n_mels

    def __len__(self):
        return len(self.file_paths)

    def __getitem__(self, idx):
        filepath = self.file_paths[idx]
        label = self.labels[idx]
        audio = load_audio(filepath, self.sr, self.max_duration)
        mel_spec = extract_mel_spectrogram(audio, self.sr, self.n_mels)
        mel_spec = torch.tensor(mel_spec, dtype=torch.float32)
        label = torch.tensor(label, dtype=torch.float32)
        return mel_spec, label


class TCNBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, dilation, dropout):
        super(TCNBlock, self).__init__()
        self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size, padding=(kernel_size - 1) * dilation,
                               dilation=dilation)
        self.relu1 = nn.ReLU()
        self.dropout1 = nn.Dropout(dropout)

        self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size, padding=(kernel_size - 1) * dilation,
                               dilation=dilation)
        self.relu2 = nn.ReLU()
        self.dropout2 = nn.Dropout(dropout)

        self.residual = nn.Conv1d(in_channels, out_channels, 1) if in_channels != out_channels else nn.Identity()

    def forward(self, x):
        out = self.conv1(x)
        out = self.relu1(out)
        out = self.dropout1(out)

        out = self.conv2(out)
        out = self.relu2(out)
        out = self.dropout2(out)

        res = self.residual(x)
        return out + res


class TCN(nn.Module):
    def __init__(self, input_size, num_classes, num_channels, kernel_size=3, dropout=0.2):
        super(TCN, self).__init__()
        layers = []
        num_blocks = len(num_channels)
        for i in range(num_blocks):
            in_channels = input_size if i == 0 else num_channels[i - 1]
            out_channels = num_channels[i]
            dilation = 2 ** i
            layers.append(TCNBlock(in_channels, out_channels, kernel_size, dilation, dropout))
        self.network = nn.Sequential(*layers)
        self.global_pool = nn.AdaptiveAvgPool1d(1)
        self.fc = nn.Linear(num_channels[-1], num_classes)

    def forward(self, x):
        x = self.network(x)
        x = self.global_pool(x).squeeze(-1)  # Global pooling
        x = self.fc(x)
        return x





def train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs=20):
    for epoch in range(num_epochs):
        model.train()
        train_loss = 0
        for mel_spec, labels in train_loader:
            mel_spec = mel_spec.to(device).permute(0, 2, 1)  # 调整维度 (batch, freq, time)
            labels = labels.to(device)

            optimizer.zero_grad()
            outputs = model(mel_spec)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            train_loss += loss.item()

        # 验证
        model.eval()
        val_loss = 0
        with torch.no_grad():
            for mel_spec, labels in val_loader:
                mel_spec = mel_spec.to(device).permute(0, 2, 1)
                labels = labels.to(device)
                outputs = model(mel_spec)
                loss = criterion(outputs, labels)
                val_loss += loss.item()

        print(
            f"Epoch [{epoch + 1}/{num_epochs}], Train Loss: {train_loss / len(train_loader):.4f}, Val Loss: {val_loss / len(val_loader):.4f}")


def evaluate_model(model, test_loader):
    model.eval()
    y_true, y_pred = [], []
    with torch.no_grad():
        for mel_spec, labels in test_loader:
            mel_spec = mel_spec.to(device).permute(0, 2, 1)
            labels = labels.cpu().numpy()
            outputs = torch.sigmoid(model(mel_spec)).cpu().numpy()

            y_true.append(labels)
            y_pred.append(outputs)

    y_true = np.vstack(y_true)
    y_pred = np.vstack(y_pred)

    # 计算平均精度（mAP）
    map_score = average_precision_score(y_true, y_pred, average='macro')
    print(f"mAP: {map_score:.4f}")

# 模型超参数
num_classes = 50  # 假设有 50 个鸟类
input_size = 128  # 梅尔频谱的频率维度
num_channels = [64, 128, 256]  # TCN 层数和通道数
kernel_size = 3
dropout = 0.2

# 创建模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = TCN(input_size, num_classes, num_channels, kernel_size, dropout).to(device)

# 损失函数和优化器
criterion = nn.BCEWithLogitsLoss()  # 二元交叉熵
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 数据路径和标签
train_files = [...]  # 训练音频文件路径列表
train_labels = [...]  # 对应的多标签列表
val_files = [...]  # 验证音频文件路径列表
val_labels = [...]
test_files = [...]
test_labels = [...]

# 数据加载器
train_dataset = BirdDataset(train_files, train_labels)
val_dataset = BirdDataset(val_files, val_labels)
test_dataset = BirdDataset(test_files, test_labels)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

# 训练
train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs=20)

# 测试
evaluate_model(model, test_loader)
